package com.zk.bytecode;

import cn.hutool.core.io.IoUtil;
import com.zk.config.UserConfiguration;
import net.bytebuddy.agent.ByteBuddyAgent;
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
import org.springframework.boot.ConfigurableBootstrapContext;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.SpringApplicationRunListener;
import org.springframework.context.ConfigurableApplicationContext;

import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.time.Duration;


public class ByteCodeSpringApplicationRunListener implements SpringApplicationRunListener {


    public ByteCodeSpringApplicationRunListener(SpringApplication application, String[] args) {
        ByteBuddyAgent.install();
        System.out.println("init ByteCodeSpringApplicationRunListener");
    }

    @Override
    public void starting(ConfigurableBootstrapContext bootstrapContext) {
        System.out.println("starting");
    }


    @Override
    public void contextPrepared(ConfigurableApplicationContext context) {
        System.out.println("contextPrepared");
    }

    @Override
    public void contextLoaded(ConfigurableApplicationContext context) {
        System.out.println("contextLoaded");
//        handleBeans(context);
    }


    @Override
    public void started(ConfigurableApplicationContext context, Duration duration) {
        System.out.println("started");
        handleBeans(context);
    }

    private void handleBeans(ConfigurableApplicationContext context) {
        DefaultListableBeanFactory beanFactory = (DefaultListableBeanFactory) context.getBeanFactory();

        Object app = beanFactory.getBean(UserConfiguration.class);
        saveClassByteCode(app.getClass());
    }

    private void saveClassByteCode(Class<?> clazz) {
        Instrumentation instrumentation = ByteBuddyAgent.getInstrumentation();
        instrumentation.addTransformer((loader, className, classBeingRedefined, protectionDomain, classfileBuffer) -> {
            if (classBeingRedefined != null) {
                try {
                    IoUtil.write(new FileOutputStream(clazz.getName() + ".class"), true, classfileBuffer);
                } catch (FileNotFoundException e) {
                    e.printStackTrace();
                }
            }

            return classfileBuffer;
        }, true);

        try {
            instrumentation.retransformClasses(clazz);
        } catch (UnmodifiableClassException e) {
            e.printStackTrace();
        }
    }


}
